# @FileName  : test-detect_text.py
# @Time      : 2025/11/8 10:44
# @Author    : LuZhaoHui
# @Software  : PyCharm

# 探测文字

import os
import cv2
import numpy as np

import cv2
import numpy as np
from PIL import Image
from PIL.ExifTags import TAGS


def get_image_orientation_exif(image_path):
    """
    通过EXIF信息获取图片方向
    """
    try:
        image = Image.open(image_path)
        exif_data = image._getexif()

        if exif_data:
            for tag_id, value in exif_data.items():
                tag = TAGS.get(tag_id, tag_id)
                if tag == 'Orientation':
                    return value
        return 1  # 默认方向
    except:
        return 1


def correct_image_orientation(image_path):
    """
    根据EXIF信息自动校正图片方向
    """
    orientation = get_image_orientation_exif(image_path)
    img = cv2.imread(image_path)

    if orientation == 3:
        # 180度旋转
        img = cv2.rotate(img, cv2.ROTATE_180)
    elif orientation == 6:
        # 90度顺时针
        img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
    elif orientation == 8:
        # 90度逆时针
        img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)

    return img


def detect_orientation_feature_matching(image):
    """
    通过特征点匹配检测图片方向
    """
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 使用ORB特征检测器
    orb = cv2.ORB_create()
    keypoints, descriptors = orb.detectAndCompute(gray, None)

    if descriptors is None:
        return "Unknown", 0

    # 旋转图像并比较特征点
    rotations = [0, 90, 180, 270]
    best_rotation = 0
    max_matches = 0

    original_kp = keypoints
    original_desc = descriptors

    for angle in rotations:
        if angle == 0:
            rotated = gray
        elif angle == 90:
            rotated = cv2.rotate(gray, cv2.ROTATE_90_CLOCKWISE)
        elif angle == 180:
            rotated = cv2.rotate(gray, cv2.ROTATE_180)
        elif angle == 270:
            rotated = cv2.rotate(gray, cv2.ROTATE_90_COUNTERCLOCKWISE)

        # 检测旋转后图像的特征点
        kp_rotated, desc_rotated = orb.detectAndCompute(rotated, None)

        if desc_rotated is not None and original_desc is not None:
            # 特征点匹配
            bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
            matches = bf.match(original_desc, desc_rotated)

            if len(matches) > max_matches:
                max_matches = len(matches)
                best_rotation = angle

    return f"{best_rotation} degrees", best_rotation


def detect_orientation_hough_lines(image):
    """
    使用Hough直线检测分析图片方向
    """
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # 检测直线
    lines = cv2.HoughLines(edges, 1, np.pi / 180, threshold=100)

    if lines is None:
        return "Unknown", 0

    # 统计角度分布
    angles = []
    for rho, theta in lines[:, 0]:
        angle = theta * 180 / np.pi
        angles.append(angle)

    # 计算主要角度
    angle_array = np.array(angles)
    main_angle = np.median(angle_array)

    # 判断方向
    if main_angle < 45 or main_angle > 135:
        return "Horizontal", 0
    else:
        return "Vertical", 90


def detect_orientation_fourier(image):
    """
    使用傅里叶变换检测图片方向
    """
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 傅里叶变换
    f = np.fft.fft2(gray)
    fshift = np.fft.fftshift(f)
    magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1)

    # 转换为极坐标
    center = (magnitude_spectrum.shape[1] // 2, magnitude_spectrum.shape[0] // 2)
    polar = cv2.linearPolar(magnitude_spectrum, center,
                            min(center[0], center[1]),
                            cv2.WARP_FILL_OUTLIERS)

    # 计算角度方向的能量
    energy_by_angle = np.sum(polar, axis=1)
    dominant_angle = np.argmax(energy_by_angle)

    return f"{dominant_angle} degrees", dominant_angle


def detect_orientation_by_text(image):
    """
    通过检测文本方向来判断图片方向
    """
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 使用OCR方法或文本检测
    # 这里使用简单的轮廓分析作为示例

    # 二值化
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # 查找轮廓
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if not contours:
        return "Unknown", 0

    # 分析轮廓方向
    horizontal_count = 0
    vertical_count = 0

    for contour in contours:
        if cv2.contourArea(contour) > 100:
            rect = cv2.minAreaRect(contour)
            width, height = rect[1]

            # 根据宽高比判断方向
            if width > height:
                horizontal_count += 1
            else:
                vertical_count += 1

    if horizontal_count > vertical_count:
        return "Horizontal", 0
    else:
        return "Vertical", 90


def detect_image_orientation_comprehensive(image_path):
    """
    综合检测图片方向
    """
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        return "无法读取图像", 0

    results = []

    # 方法1: EXIF信息
    # try:
    #     orientation_exif = get_image_orientation_exif(image_path)
    #     if orientation_exif != 1:
    #         results.append(("EXIF", orientation_exif))
    # except:
    #     pass

    # 方法2: 特征点匹配
    # orientation_feature, angle_feature = detect_orientation_feature_matching(image)
    # if orientation_feature != "Unknown":
    #     results.append(("Feature", angle_feature))

    # 方法3: 直线检测
    # orientation_line, angle_line = detect_orientation_hough_lines(image)
    # if orientation_line != "Unknown":
    #     results.append(("Hough", angle_line))

    # 方法4: 文本方向
    orientation_text, angle_text = detect_orientation_by_text(image)
    if orientation_text != "Unknown":
        results.append(("Text", angle_text))

    # 投票决定最终方向
    if not results:
        return "Unknown", 0

    # 统计最常见的角度
    angles = [angle for _, angle in results]
    final_angle = max(set(angles), key=angles.count)

    angle_names = {
        0: "0° (正常)",
        90: "90° (顺时针)",
        180: "180° (倒立)",
        270: "270° (逆时针)"
    }

    return angle_names.get(final_angle, f"{final_angle}°"), final_angle


def correct_image_orientation_by_angle(image_path, angle):
    """
    根据角度校正图片方向
    """
    img = cv2.imread(image_path)

    if angle == 90:
        corrected = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
    elif angle == 180:
        corrected = cv2.rotate(img, cv2.ROTATE_180)
    elif angle == 270:
        corrected = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
    else:
        corrected = img

    return corrected


def main():
    # 使用示例
    dir = 'f:/image/haidian/test'
    for f in os.listdir(dir):
        image_path = os.path.join(dir, f)
        # 综合检测
        orientation, angle = detect_image_orientation_comprehensive(image_path)
        print("%s 图片方向: %s" % (image_path, orientation))

        # 自动校正
        # if angle != 0:
        #     corrected_image = correct_image_orientation_by_angle(image_path, angle)
        #     cv2.imshow("Corrected Image", corrected_image)
        #     cv2.waitKey(0)
        #     cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
